from dimredu.eRPCAviaADMMFast import eRPCA as eRPCASparse
import numpy as np

import matplotlib.pyplot as plt
from sklearn.datasets import fetch_lfw_people


def denseToSparse(M, E):
    assert M.shape == E.shape, 'shape mismatch'
    m = M.shape[0]
    n = M.shape[1]

    u = np.empty([m * n])
    v = np.empty([m * n])
    vecM = np.empty([m * n])
    vecE = np.empty([m * n])

    k = 0
    for i in range(m):
        for j in range(n):
            u[k] = i
            v[k] = j
            vecM[k] = M[i, j]
            vecE[k] = E[i, j]
            k += 1

    return m, n, u, v, vecM, vecE


def eRPCA(M, E, **kw):
    m, n, u, v, vecM, vecE = denseToSparse(M, E)
    maxRank = np.min(M.shape)
    return eRPCASparse(m, n, u, v, vecM, vecE, maxRank, **kw)


def plot_faces(faces):
    fig, axes = plt.subplots(6, 6, figsize=(10, 10),
                              subplot_kw={'xticks': [], 'yticks': []},
                              gridspec_kw=dict (hspace=0.1, wspace=0.1))
    for i, ax in enumerate (axes.flat):
        ax.imshow(faces[i].reshape(62, 47), cmap='bone')
    plt.show()


def test_small():
    X = np.random.random(size=[5, 15])
    E = np.ones(X.shape)*1e-6
    eRPCA(X, E)


if __name__ == '__main__':
    faces = fetch_lfw_people()
    random_indexes = np.random.permutation(len(faces.data))
    X = faces.data[random_indexes]
    example_faces = X[:36, :]
    # plot_faces(example_faces)


    E = np.ones(example_faces.shape) * 1e-6
    eRPCA(example_faces, E)
    # result = eRPCA(example_faces, E)

    # test_small()
    # U = result[0]
    # E = result[1]
    # VT = result[2]
    # S = result[3]
    # low_Rank_Matrix = np.dot(U, E, VT)
    # sparse_Matrix = S
    # plot_faces(low_Rank_Matrix)
    # plot_faces(sparse_Matrix)